package com.shujia.flink.kafka;

import com.alibaba.fastjson2.JSON;
import com.alibaba.fastjson2.annotation.JSONField;
import com.shujia.flink.util.DateUtils;
import lombok.AllArgsConstructor;
import lombok.Data;
import org.apache.flink.api.common.eventtime.WatermarkStrategy;
import org.apache.flink.api.common.serialization.SimpleStringSchema;
import org.apache.flink.api.common.typeinfo.TypeInformation;
import org.apache.flink.configuration.Configuration;
import org.apache.flink.connector.kafka.source.KafkaSource;
import org.apache.flink.connector.kafka.source.enumerator.initializer.OffsetsInitializer;
import org.apache.flink.streaming.api.datastream.DataStream;
import org.apache.flink.streaming.api.datastream.KeyedStream;
import org.apache.flink.streaming.api.datastream.WindowedStream;
import org.apache.flink.streaming.api.environment.StreamExecutionEnvironment;
import org.apache.flink.streaming.api.functions.sink.RichSinkFunction;
import org.apache.flink.streaming.api.functions.windowing.ProcessWindowFunction;
import org.apache.flink.streaming.api.windowing.assigners.SlidingEventTimeWindows;
import org.apache.flink.streaming.api.windowing.time.Time;
import org.apache.flink.streaming.api.windowing.windows.TimeWindow;
import org.apache.flink.util.Collector;

import java.sql.Connection;
import java.sql.DriverManager;
import java.sql.PreparedStatement;
import java.time.Duration;

/**
 * @author shujia
 */
public class Demo3FlumeCar {
    public static void main(String[] args) throws Exception {
        //创建环境
        StreamExecutionEnvironment env = StreamExecutionEnvironment.getExecutionEnvironment();

        env.setParallelism(1);

        //创建kafka source
        KafkaSource<String> source = KafkaSource.<String>builder()
                .setBootstrapServers("master:9092,node1:9092,node2:9092")
                .setTopics("flume_car")
                .setGroupId("Demo3FlumeCar")
                .setStartingOffsets(OffsetsInitializer.latest())
                .setValueOnlyDeserializer(new SimpleStringSchema())
                .build();

        //从kafka中读取数据
        DataStream<String> kafkaSource = env
                .fromSource(source, WatermarkStrategy.noWatermarks(), "kafka source");

        //解析数据
        DataStream<Car> cars = kafkaSource.map(line -> JSON.parseObject(line, Car.class), TypeInformation.of(Car.class));

        //指定时间字段和水位线
        WatermarkStrategy<Car> watermarkStrategy = WatermarkStrategy
                //指定水位线
                .<Car>forBoundedOutOfOrderness(Duration.ofSeconds(5))
                //指定时间字段，使用时间戳格式
                .withTimestampAssigner((car, t) -> DateUtils.stringToTimestamp(car.getTimestamp()));
        //指定时间字段和水位线
        DataStream<Car> carsAss = cars.assignTimestampsAndWatermarks(watermarkStrategy);

        /*
         * 实时统计道路拥堵情况
         * 1、计算最近一段时间道路的平均车速和车流量
         * 2、计算最近15秒，。每隔5秒计算一次
         */
        KeyedStream<Car, String> keyBys = carsAss.keyBy(Car::getRoad);

        //划分窗口
        WindowedStream<Car, String, TimeWindow> windows = keyBys.window(SlidingEventTimeWindows.of(Time.seconds(15), Time.seconds(5)));

        //计算平均车速和车流量
        DataStream<RoadResult> results = windows
                .process(new ProcessWindowFunction<Car, RoadResult, String, TimeWindow>() {
                    /**
                     * process方法每一个窗口执行一次
                     * @param road 道路
                     * @param context 上下文对象
                     * @param elements 窗口内的数据
                     * @param out 用于将数据发生到下游
                     */
                    @Override
                    public void process(String road,
                                        ProcessWindowFunction<Car, RoadResult, String, TimeWindow>.Context context,
                                        Iterable<Car> elements,
                                        Collector<RoadResult> out) {
                        long flow = 0;
                        Double sumSpeed = 0.0;
                        for (Car car : elements) {
                            flow++;
                            sumSpeed += car.getSpeed();
                        }
                        //计算平均车速
                        Double avgSpeed = sumSpeed / flow;

                        //获取窗口的时间
                        String start = DateUtils.timestampToString(context.window().getStart());
                        String end = DateUtils.timestampToString(context.window().getEnd());


                        //将数据发生到下游
                        out.collect(new RoadResult(road, start, end, flow, avgSpeed));
                    }
                });


        //将结果保存到数据库
        results.addSink(new RichSinkFunction<RoadResult>() {
            private Connection con;
            private PreparedStatement stat;

            @Override
            public void open(Configuration parameters) throws Exception {
                //获取数据库连接
                Class.forName("com.mysql.jdbc.Driver");
                con = DriverManager.getConnection("jdbc:mysql://master:3306/student", "root", "123456");
                stat = con.prepareStatement("insert into road_result values(?,?,?,?,?)");
            }

            @Override
            public void close() throws Exception {
                //关闭连接
                con.close();
                stat.close();
            }

            @Override
            public void invoke(RoadResult roadResult, Context context) throws Exception {
                //写入数据
                stat.setString(1, roadResult.getRoad());
                stat.setString(2, roadResult.getStart());
                stat.setString(3, roadResult.getEnd());
                stat.setLong(4, roadResult.getFlow());
                stat.setDouble(5, roadResult.getAvgSpeed());

                stat.execute();
            }
        });

        env.execute();


    }

    @Data
    @AllArgsConstructor
    static class Car {
        @JSONField(name = "tollgate_id")
        private String tollgateId;
        @JSONField(name = "license_plate")
        private String licensePlate;
        private String timestamp;
        private Double longitude;
        private Double latitude;
        private Integer speed;
        private String city;
        private String road;
        private String district;
    }

    @Data
    @AllArgsConstructor
    static class RoadResult {
        private String road;
        private String start;
        private String end;
        private long flow;
        private Double avgSpeed;

    }
}
